# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: jaewon, inkyu

import datetime
import logging
import math
import sys
import time

import tabulate

from coin.base.param_util import to_list
from coin.exchange.base.order_gateway import OrderSide, OrderType
from coin.exchange.bitfinex_v2.kr_rest.product import BitfinexProduct
from coin.util.basis_receiver import BasisReceiver

from coin.strategy.mm.calculator import MovingSum
from coin.exchange.bitmex.kr_rest.futures_product import BitmexFuturesProduct
from coin.strategy.mm.bitfinex.simple_mm_executor import BitfinexSimpleMMExecutor
from coin.strategy.mm.base.bitmex_bitfinex_strategy_base import StrategyBase
from coin.exchange.shared.kr_rest.product import SharedProduct


def ceil_price(price, tick_size):
  return math.ceil(price / tick_size) * tick_size


def floor_price(price, tick_size):
  return math.floor(price / tick_size) * tick_size


def norm_float(value, precision):
  return float(('%%.%df' % precision) % value)


class SoftArbStrategy(StrategyBase):
  def __init__(self, strategy_params, dry_run=True):
    super().__init__(product_1=strategy_params['product_agg'],
                     product_2=strategy_params['product_passive'],
                     balance_currency=strategy_params.get('balance_currency', None),
                     logger=logging.getLogger('SoftArbStrategy'))

    self._strategy_params = strategy_params
    self._dry_run = dry_run

    self._ask_diff = 0
    self._bid_diff = 0
    self._ask_basis = None
    self._bid_basis = None
    self._edge = strategy_params['edge_bp'] / 10000.

    self._true_book_2 = strategy_params['passive_true_book']
    self._lot_size_usd_2 = strategy_params['lot_size']
    self._tick_size_2 = strategy_params['tick_size']
    self._base_pos_2 = strategy_params['passive_base_pos']

    self._agg_edge = strategy_params['agg_edge_bp'] / 10000.
    self._true_book_1 = strategy_params['agg_true_book']
    self._lot_size_usd_min_1 = strategy_params['agg_lot_size_min']
    self._lot_size_usd_max_1 = strategy_params['agg_lot_size_max']
    self._qty_precision_1 = strategy_params['qty_precision']
    self._base_pos_1 = strategy_params['agg_base_pos']

    self._passive_executor = None
    self._basis_receiver = None

  def on_base_init(self, ioloop, queue, flow_sub):
    super().on_base_init(ioloop, queue, flow_sub)
    self._basis_receiver = BasisReceiver(ioloop)
    self._basis_receiver.start()
    self.start_print_feature_status()

  def init_executor(self):
    # Intialize passive executor
    self._logger.info('Intializing passive executor...')
    passive_executor_params = {'product': self._strategy_params['product_passive']}
    for field in ['lot_size', 'stickiness', 'min_pos', 'max_pos']:
      passive_executor_params[field] = self._strategy_params[field]
    self._passive_executor = BitfinexSimpleMMExecutor(passive_executor_params,
                                                      self._og_2,
                                                      self._logger)

    # Register fill callback
    self._og_2.register_fill_callback(self.on_passive_fill)

  def get_basis(self):
    product = SharedProduct.FromProduct(self._product_2)
    ask_basis, ask_basis_time = self._basis_receiver.get_basis('ask', 'BITFINEX', 'BITMEX', product,
                                                               36000)
    bid_basis, bid_basis_time = self._basis_receiver.get_basis('bid', 'BITFINEX', 'BITMEX', product,
                                                               36000)
    if ask_basis is not None and bid_basis is not None:
      if time.time() - ask_basis_time > 10:  # 10 sec
        self._logger.warn('Too old ask basis data')
        return (None, None)
      if time.time() - bid_basis_time > 10:  # 10 sec
        self._logger.warn('Too old bid basis data')
        return (None, None)
    else:
      self._logger.warn('ask_basis: %s, bid_basis: %s' % (ask_basis, bid_basis))
      return (None, None)
    return (-ask_basis, -bid_basis)

  def start_print_feature_status(self):
    try:
      product = SharedProduct.FromProduct(self._product_2)
      ask_basis, ask_basis_time = self._basis_receiver.get_basis('ask', 'BITFINEX', 'BITMEX',
                                                                 product, 36000)
      bid_basis, bid_basis_time = self._basis_receiver.get_basis('bid', 'BITFINEX', 'BITMEX',
                                                                 product, 36000)
      ask_basis10, ask_basis_time = self._basis_receiver.get_basis('ask', 'BITFINEX', 'BITMEX',
                                                                   product, 10)
      bid_basis10, bid_basis_time = self._basis_receiver.get_basis('bid', 'BITFINEX', 'BITMEX',
                                                                   product, 10)
      ask_basis600, ask_basis_time = self._basis_receiver.get_basis('ask', 'BITFINEX', 'BITMEX',
                                                                    product, 600)
      bid_basis600, bid_basis_time = self._basis_receiver.get_basis('bid', 'BITFINEX', 'BITMEX',
                                                                    product, 600)

      if ask_basis is not None and bid_basis is not None:
        if time.time() - ask_basis_time > 10:  # 10 sec
          self._logger.warn('Too old ask basis data')
          return
        if time.time() - bid_basis_time > 10:  # 10 sec
          self._logger.warn('Too old bid basis data')
          return
      else:
        self._logger.warn('ask_basis: %s, bid_basis: %s' % (ask_basis, bid_basis))
        return
      self._logger.info('BASIS %s: %s %s %s %s %s %s %s %s' % (self._product_2,
                                                               self._ask_diff,
                                                               self._bid_diff,
                                                               -ask_basis,
                                                               -bid_basis,
                                                               -ask_basis10,
                                                               -bid_basis10,
                                                               -ask_basis600,
                                                               -bid_basis600))
    finally:
      self._ioloop.add_timeout(datetime.timedelta(seconds=10), self.start_print_feature_status)

  def _print_imbalance(self):
    if not self.is_og_ready() or self._book_1 is None or self._book_2 is None:
      return

    mid_price_1 = (self._book_1.ask0().price + self._book_1.bid0().price) / 2.
    balance_1 = ((self._og_1.get_balance('BTC') - self._base_pos_1) * mid_price_1
                 + self._og_1.get_position('BTC-USD.PERPETUAL'))

    mid_price_2 = (self._book_2.ask0().price + self._book_2.bid0().price) / 2.
    balance_2 = (self._og_2.get_balance('BTC') - self._base_pos_2) * mid_price_2

    self._logger.info('Imbalance %.1f USD = Bitmex:%.1f + Bitfinex:%.1f',
                      balance_1 + balance_2,
                      balance_1,
                      balance_2)

  def _print_balance(self):  # Override the original version.
    if self._book_1 is None or self._book_2 is None:
      return

    balance_1 = self._og_1.get_balances_copy()
    balance_2 = self._og_2.get_balances_copy()

    # Adjust bitmex balance
    mid_price_1 = (self._book_1.ask0().price + self._book_1.bid0().price) / 2.
    pos_usd_1 = self._og_1.get_position('BTC-USD.PERPETUAL')
    pos_1 = pos_usd_1 / mid_price_1
    balance_1['BTC'] = balance_1.get('BTC', 0) + pos_1
    balance_1['USDT'] = -pos_usd_1

    if self._balance_currency is not None:
      balance_currency = self._balance_currency
    else:
      balance_currency = sorted(set(balance_1.keys() + balance_2.keys()))

    table = []
    for cur in balance_currency:
      bal_1 = balance_1.get(cur, 0.)
      bal_2 = balance_2.get(cur, 0.)
      table.append(['*', cur, bal_1 + bal_2, bal_1, bal_2])

    self._logger.info(
        '----- Balance:\n* %s\n%s' %
        (str(datetime.datetime.now()),
         tabulate.tabulate(
             table, headers=['*', 'Currency', 'Total', 'Bitmex', 'Bitfinex'], floatfmt=".6f")))

    self._print_imbalance()

  def on_book(self):
    if self._book_1 is None or self._book_2 is None:
      return

    # Check book freshness
    book_ts = max(self._book_1.timestamp, self._book_2.timestamp)
    min_ts = min(self._book_1.timestamp, self._book_2.timestamp)
    if book_ts - min_ts > (2 * 10**9):
      return

    # Check book sanity
    if (self._book_1.ask0().price <= self._book_1.bid0().price
        or self._book_2.ask0().price <= self._book_2.bid0().price):
      return

    askt_1 = self._book_1.get_notional_asks_by_qty(self._true_book_1)[1][-1][0]
    bidt_1 = self._book_1.get_notional_bids_by_qty(self._true_book_1)[1][-1][0]

    ask0_2 = self._book_2.ask0().price
    bid0_2 = self._book_2.bid0().price
    askt_2 = self._book_2.get_notional_asks_by_amt(self._true_book_2)[1][-1][0]
    bidt_2 = self._book_2.get_notional_bids_by_amt(self._true_book_2)[1][-1][0]

    self._ask_diff = askt_1 - askt_2
    self._bid_diff = bidt_1 - bidt_2

    ask_basis, bid_basis = self.get_basis()
    if not ask_basis or not bid_basis:
      return
    self._ask_basis = ask_basis
    self._bid_basis = bid_basis

    askt_1_proj_2 = askt_1 - float(ask_basis)
    bidt_1_proj_2 = bidt_1 - float(bid_basis)

    askp_2 = max(
        askt_1_proj_2 * (1. + self._edge),
        #                 ask0_2 - self._tick_size_2,
        bid0_2 + self._tick_size_2)
    bidp_2 = min(
        bidt_1_proj_2 * (1. - self._edge),
        #                 bid0_2 + self._tick_size_2,
        ask0_2 - self._tick_size_2)
    askp_2 = ceil_price(askp_2, self._tick_size_2)
    bidp_2 = floor_price(bidp_2, self._tick_size_2)

    if not self.is_og_ready():
      return

    if self._dry_run:
      print("dry run", time.time())
      return

    fired = self._passive_executor.manage_mm_orders(askp_2, bidp_2)

    if fired:
      # If there is no new passive order fired, do not cancel aggression orders.
      self.cancel_all_aggression_orders()
      self._logger.info('[2] SELL %f @ %.1f (%.1f), BUY %f @ %.1f (%.1f)',
                        self._lot_size_usd_2,
                        askp_2,
                        askp_2 - ask0_2,
                        self._lot_size_usd_2,
                        bidp_2,
                        bid0_2 - bidp_2)

  def on_passive_fill(self, fill_price, fill_qty, order):
    if (order.product != self._product_2 or self._dry_run):
      return

    if not self.is_og_ready() or self._dry_run:
      self._logger.error('OG is not ready. Skip aggression')
      return

    if order.order_side == OrderSide.SELL:
      self.try_aggression_buy(fill_price, fill_qty)
    elif order.order_side == OrderSide.BUY:
      self.try_aggression_sell(fill_price, fill_qty)
    else:
      raise ValueError('Invalid side: %s' % str(order.order_side))

  def try_aggression_buy(self, fill_price, fill_qty):
    if self._dry_run or not self._ask_basis:
      return

    askt_1 = self._book_1.get_notional_asks_by_qty(self._true_book_1)[1][-1][0]
    askt_1_proj_2 = askt_1 - float(self._ask_basis)
    if (fill_price / askt_1_proj_2) < (1. + self._agg_edge):
      self._logger.error(
          '[1] SELL-BUY: Not enough edge: '
          'fill=%.1f, askt_1=%.4f, askt_1_proj_2=%.4f, edge=%.2fbp',
          fill_price,
          askt_1,
          askt_1_proj_2, (fill_price / askt_1_proj_2 - 1.) * 10000.0)
      return

    # TODO(jaewon): Parameterize 'btc' and 'XBTUSD'
    pos_2 = self._og_2.get_balance('BTC')

    mid_price_1 = (self._book_1.ask0().price + self._book_1.bid0().price) / 2.
    pos_usd_1 = self._og_1.get_position('BTC-USD.PERPETUAL')
    pos_usd_1 -= fill_qty
    pos_1 = self._og_1.get_balance('BTC') + (pos_usd_1 / mid_price_1)

    lot_size_min = self._lot_size_usd_min_1
    lot_size_max = self._lot_size_usd_max_1

    print(self._base_pos_1, self._base_pos_2, pos_1, pos_2)
    diff = ((self._base_pos_1 + self._base_pos_2) - (pos_1 + pos_2)) * mid_price_1
    if diff < lot_size_min:
      self._logger.info('[1] NO BUY: qty=%.4f < lot_size_min=%.4f', diff, self._lot_size_usd_min_1)
      return

    buy_qty = norm_float(min(diff, lot_size_max), self._qty_precision_1)
    self._og_1.submit(self._product_1.native_symbol,
                      "%.5f" % (askt_1),
                      "%.5f" % (buy_qty),
                      OrderSide.BUY,
                      OrderType.LIMIT)
    self._logger.info('[1] BUY  %.4f @ %.4f', buy_qty, askt_1)

  def try_aggression_sell(self, fill_price, fill_qty):
    if self._dry_run or not self._bid_basis:
      return

    bidt_1 = self._book_1.get_notional_bids_by_qty(self._true_book_1)[1][-1][0]
    bidt_1_proj_2 = bidt_1 - float(self._bid_basis)
    if (bidt_1_proj_2 / fill_price) < (1. + self._agg_edge):
      self._logger.error(
          '[1] BUY-SELL: Not enough edge: '
          'fill=%.1f, bidt_1=%.4f, bidt_1_proj_2=%.4f, edge=%.2fbp',
          fill_price,
          bidt_1,
          bidt_1_proj_2, (bidt_1_proj_2 / fill_price - 1.) * 10000.0)
      return

    # TODO(jaewon): Parameterize 'btc' and 'XBTUSD'
    pos_2 = self._og_2.get_balance('BTC')

    mid_price_1 = (self._book_1.ask0().price + self._book_1.bid0().price) / 2.
    pos_usd_1 = self._og_1.get_position('BTC-USD.PERPETUAL')
    pos_usd_1 -= fill_qty
    pos_1 = self._og_1.get_balance('BTC') + (pos_usd_1 / mid_price_1)

    lot_size_min = self._lot_size_usd_min_1
    lot_size_max = self._lot_size_usd_max_1

    print(self._base_pos_1, self._base_pos_2, pos_1, pos_2)

    diff = ((self._base_pos_1 + self._base_pos_2) - (pos_1 + pos_2)) * mid_price_1
    if -diff < lot_size_min:
      self._logger.info('[1] NO SELL: qty=%.4f < lot_size_min=%.4f',
                        -diff,
                        self._lot_size_usd_min_1)
      return

    sell_qty = norm_float(min(-diff, lot_size_max), self._qty_precision_1)
    self._og_1.submit(self._product_1.native_symbol,
                      "%.5f" % (bidt_1),
                      "%.5f" % (sell_qty),
                      OrderSide.SELL,
                      OrderType.LIMIT)
    self._logger.info('[1] SELL %.4f @ %.4f', sell_qty, bidt_1)

  def cancel_all_aggression_orders(self):
    if self._dry_run:
      return
    if self._og_1.get_working_order():
      self._og_1.cancel_all()


def main(argv):
  # Strategy parameter
  product_passive = BitfinexProduct.FromStr('BTC-USDT')
  product_agg = BitmexFuturesProduct.FromStr('BTC-USD.PERPETUAL')
  lot_size = 0.1  # BTC
  lot_size_usd = lot_size * 10000  # USD, temp
  stack = 10

  passive_base_pos = 1.5  # Base currency
  agg_base_pos = 2.0137

  strategy_params = {
      'edge_bp': 15.,
      'basis_ma_window_min': 60,

  # Passive-posting parameters
      'product_passive': product_passive,
      'passive_true_book': 1000.,
      'lot_size': lot_size,
      'stickiness': 0.,
      'min_pos': passive_base_pos - lot_size * stack,
      'max_pos': passive_base_pos + lot_size * stack,
      'tick_size': 1,
      'passive_base_pos': passive_base_pos,

  # Aggression parameters
      'product_agg': product_agg,
      'agg_true_book': 1000.,
      'agg_lot_size_min': 1,  # USD
      'agg_lot_size_max': 3 * lot_size_usd,  # USD
      'agg_edge_bp': -5.,
      'qty_precision': 0,  # USD
      'agg_base_pos': agg_base_pos,

  # Other parameters
      'balance_currency': ['BTC', 'USDT'],
  }

  # Trading environment config
  exit_after_min = 30
  env_config = {
      'first_product': product_passive,
      'first_og_config_name': 'bitfinex_jshin0',
      'second_product': product_agg,
      'second_og_config_name': 'bitmex_jshin4',
      'init_order_gateway_after': datetime.timedelta(minutes=0),
      'exit_after': datetime.timedelta(minutes=exit_after_min)
  }

  # Run strategy!
  strategy = SoftArbStrategy(strategy_params, dry_run=False)
  return StrategyBase.run_live(on_base_init_callback=strategy.on_base_init,
                               on_book_reset_callback=strategy.on_book_reset,
                               on_order_gateway_reset_callback=strategy.on_og_reset,
                               **env_config)


if __name__ == '__main__':
  logging.basicConfig(level='DEBUG', format='%(levelname)8s %(asctime)s %(name)s] %(message)s')
  sys.exit(main(sys.argv))
